#include <math.h>

struct FfiVector {
    float *data;
    int len;
};

__global__ void vector_add(float *x, float *y, float *z, size_t n) {
    int idx = threadIdx.x + blockDim.x * blockIdx.x;

    if (idx < n) {
        z[idx] = x[idx] + y[idx];
    }
}

extern "C" {
void vector_add_wrapper(FfiVector *x, FfiVector *y, FfiVector *z) {
    float *d_x;
    float *d_y;
    float *d_z;
    size_t vec_size = x->len * sizeof(float);

    cudaMalloc(&d_x, vec_size);
    cudaMalloc(&d_y, vec_size);
    cudaMalloc(&d_z, vec_size);

	cudaMemcpy(d_x, x->data, vec_size, cudaMemcpyHostToDevice);
	cudaMemcpy(d_y, y->data, vec_size, cudaMemcpyHostToDevice);
	
	const int num_threads = 32;
	const int num_blocks = ceil((float)x->len / (float)num_threads);
	
	vector_add<<<num_blocks, num_threads>>>(d_x, d_y, d_z, x->len);
	
	cudaMemcpy(z->data, d_z, vec_size, cudaMemcpyDeviceToHost);
	
	cudaFree(d_x);
	cudaFree(d_y);
	cudaFree(d_z);
}
}